package com.google.code.scm;

import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;
import com.sun.net.httpserver.HttpServer;
import org.apache.cassandra.db.ColumnFamilyStoreMBean;
import org.apache.cassandra.db.compaction.CompactionInfo;
import org.apache.cassandra.db.compaction.CompactionManager;
import org.apache.cassandra.db.compaction.CompactionManagerMBean;

import org.apache.cassandra.service.StorageServiceMBean;

import javax.management.JMX;
import javax.management.MBeanServerConnection;
import javax.management.MalformedObjectNameException;
import javax.management.ObjectName;
import javax.management.remote.JMXConnector;
import javax.management.remote.JMXConnectorFactory;
import javax.management.remote.JMXServiceURL;
import java.io.IOException;
import java.io.OutputStream;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryUsage;
import java.net.InetSocketAddress;
import java.text.NumberFormat;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @author smeet, Hector Izquierdo
 */
public class CassandraMonitoring {

    private static final String FMT_URL = "service:jmx:rmi:///jndi/rmi://%s:%d/jmxrmi";
    private static final String SS_OBJECT_NAME = "org.apache.cassandra.db:type=StorageService";
    private static final String CF_OBJECT_NAME = "org.apache.cassandra.db:type=ColumnFamilies,keyspace=%s,columnfamily=%s";


    private static final int DEFAULT_PORT = 8080;

    private int port = DEFAULT_PORT;

    private String host;


    public CassandraMonitoring(String host) {
        this.host = host;
    }


    /**
     * Create a connection to the JMX agent and setup the M[X]Bean proxies.
     *
     * @param kscfs KeySpace:ColumnFamily pairs
     * @throws java.io.IOException on connection failures
     */
    public void start(String[] kscfs) throws IOException {
        JMXServiceURL jmxUrl = new JMXServiceURL(String.format(FMT_URL, host, port));
        JMXConnector jmxc = JMXConnectorFactory.connect(jmxUrl, null);
        MBeanServerConnection mbeanServerConn = jmxc.getMBeanServerConnection();

        HttpServer httpServer = HttpServer.create();

        try {

            httpServer.createContext("/nodeinfo", new NHandler(
                    JMX.newMBeanProxy(mbeanServerConn, new ObjectName(SS_OBJECT_NAME), StorageServiceMBean.class),
                    ManagementFactory.newPlatformMXBeanProxy(mbeanServerConn, ManagementFactory.MEMORY_MXBEAN_NAME, MemoryMXBean.class),
                    JMX.newMBeanProxy(mbeanServerConn, new ObjectName(CompactionManager.MBEAN_OBJECT_NAME), CompactionManagerMBean.class)
            ));

            for (String kscf : kscfs) {
                String split[] = kscf.split(":");
                httpServer.createContext("/" + split[0] + "/" + split[1], new CFHandler(JMX.newMBeanProxy(mbeanServerConn, new ObjectName(String.format(CF_OBJECT_NAME, split[0], split[1])), ColumnFamilyStoreMBean.class)));
            }

        } catch (MalformedObjectNameException e) {
            throw new RuntimeException("Cannot connect.", e);
        }

        httpServer.bind(new InetSocketAddress(9090), 0);
        httpServer.start();
    }

    public static class NHandler implements HttpHandler {

        private StorageServiceMBean _ssProxy;
        private MemoryMXBean _memProxy;
        private CompactionManagerMBean _mcmProxy;

        public NHandler(StorageServiceMBean _ssProxy, MemoryMXBean _memProxy, CompactionManagerMBean _mcmProxy) {
            this._mcmProxy = _mcmProxy;
            this._memProxy = _memProxy;
            this._ssProxy = _ssProxy;
        }


        public void handle(HttpExchange httpExchange) throws IOException {
            try {
                httpExchange.getResponseHeaders().put("Content-Type", Collections.singletonList("text/plain; charset=UTF8"));
                StringBuilder body = new StringBuilder();

                // Storage Service
                append(body, "OperationMode", _ssProxy.getOperationMode());
                double load = _ssProxy.getLoad() / (1024 * 1024 * 1024);
                NumberFormat formatter = NumberFormat.getInstance();
                formatter.setGroupingUsed(false);
                append(body, "Load", formatter.format(load));


                // Compactions

                List<CompactionInfo> compactions = _mcmProxy.getCompactions();
                for (CompactionInfo compaction : compactions) {

                    append(body, "ColumnFamilyInProgress", compaction.getColumnFamily());
                    Long bytesCompacted = compaction.getBytesComplete();
                    append(body, "BytesCompacted", bytesCompacted != null ? bytesCompacted : 0);
                    Long bytesTotalInProgress = compaction.getTotalBytes();
                    append(body, "BytesTotalInProgress", bytesTotalInProgress != null ? bytesTotalInProgress : 0);

                    append(body, "CompactionType", compaction.getTaskType().toString());
                }

                MemoryUsage heapMemoryUsage = _memProxy.getHeapMemoryUsage();
                append(body, "HeapUsed", heapMemoryUsage.getUsed());

                OutputStream out = httpExchange.getResponseBody();

                byte[] bytes = body.toString().getBytes("UTF8");
                httpExchange.sendResponseHeaders(200, bytes.length);
                out.write(bytes);
                httpExchange.close();


            } catch (Throwable th) {
                th.printStackTrace();
            }
        }
    }

    public static class CFHandler implements HttpHandler {

        private ColumnFamilyStoreMBean _spProxy;

        public CFHandler(ColumnFamilyStoreMBean _spProxy) {
            this._spProxy = _spProxy;
        }


        public void handle(HttpExchange httpExchange) throws IOException {
            try {


                httpExchange.getResponseHeaders().put("Content-Type", Collections.singletonList("text/plain; charset=UTF8"));
                StringBuilder body = new StringBuilder();


                // Storage Proxy
                append(body, "ReadOperations", _spProxy.getReadCount());
                append(body, "WriteOperations", _spProxy.getWriteCount());
                append(body, "RecentReadLatencyMillis", _spProxy.getRecentReadLatencyMicros() / 1000);
                append(body, "RecentWriteLatencyMillis", _spProxy.getRecentWriteLatencyMicros() / 1000);
                append(body, "ReadLatencyMillis", _spProxy.getRecentReadLatencyMicros() / 1000);
                append(body, "WriteLatencyMillis", _spProxy.getRecentWriteLatencyMicros() / 1000);
                append(body, "LiveDiskSpaceUsed", _spProxy.getLiveDiskSpaceUsed());
                append(body, "LiveSSTableCount", _spProxy.getLiveSSTableCount());
                append(body, "PendingTasks", _spProxy.getPendingTasks());
                append(body, "RecentBloomFilterFalsePositiveRatio", _spProxy.getRecentBloomFilterFalseRatio());
                append(body, "BloomFilterFalsePositiveRatio", _spProxy.getBloomFilterFalseRatio());
                append(body, "MemtableColumnCount", _spProxy.getMemtableColumnsCount());

                long histogram[] = _spProxy.getRecentSSTablesPerReadHistogram();
                long totalReads = 0;
                for (int i = 0; i < histogram.length; i++) {
                    append(body, "SSTablesPerRead" + (i + 1), histogram[i]);
                    totalReads += histogram[i];
                }

                append(body, "RecentReadOperations", totalReads);


                OutputStream out = httpExchange.getResponseBody();

                byte[] bytes = body.toString().getBytes("UTF8");
                httpExchange.sendResponseHeaders(200, bytes.length);
                out.write(bytes);
                httpExchange.close();

            } catch (Throwable th) {
                th.printStackTrace();
            }
        }


    }

    private static void append(StringBuilder body, String name, Object value) {
        body.append(name).append('=').append(value).append('\n');
    }

    public static void main(String[] args) throws IOException, InterruptedException {
        String host = args[0];
        CassandraMonitoring server = new CassandraMonitoring(host);
        server.start(args[1].split(","));
    }
}